/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "uhm/common.hxx"
#include "uhm/const.hxx"
#include "uhm/util.hxx"

#include "uhm/object.hxx"

#include "uhm/operation/scheduler.hxx"
#include "uhm/operation/element.hxx"

#include "uhm/mesh/node.hxx"
#include "uhm/mesh/element.hxx"

#include "uhm/matrix/uhm/matrix.hxx"

#include "uhm/mesh/mesh.hxx"

#include "uhm/matrix/uhm/fla.hxx"

namespace uhm {
  // --------------------------------------------------------------
  // ** Mesh
  void Mesh_::set_rhs() {
    // copy b to x for leaf. clear for non-leaf
    std::map< int, Element_ >::iterator eit;
    for (eit=this->elements.begin();eit!=this->elements.end();eit++) {
      Element e = &(eit->second);
      assert(e->is_matrix_created());
      e->get_matrix()->set_rhs( e->is_leaf() );
    }
  }

  void Mesh_::create_matrix_without_buffer(int datatype, int n_rhs) {
    int sym = this->get_symmetry();

    std::map< int, Element_ >::iterator it;
      
    for (it=this->elements.begin();it!=this->elements.end();++it) {
      Element e = &(it->second);
        
      if (e->is_matrix_reusable()) {
        assert(e->is_matrix_created()); 
      } else {
        // create matrix
        // if the matrix is not reusable, which means connectivity is 
        // changed by p or h refinement, delete matrix object and 
        // create new one, which fit to new connectivity
        std::pair<int,int> n_dof = e->get_n_dof();
          
        // TODO :: replace the routine, fix the matrix object
        Matrix hm = new Matrix_FLA_(datatype, 
                                    n_dof.first, n_dof.second, 
                                    n_rhs, sym );
        hm->create_without_buffer();
          
        // if matrix already created, delete that object first
        if (e->is_matrix_created()) delete e->get_matrix();
        e->set_matrix(hm);
      }
    }
  }

  void Mesh_::create_leaf_matrix_buffer() {
    assert(this->get_scheduler()->is_loaded());
    Scheduler s = this->get_scheduler();
    s->execute_leaves_seq( &(op_create_matrix_buffer_with_schur) );
  }
  
  void Mesh_::create_element_matrix_buffer() { 
    this->create_element_matrix_buffer(false); 
  }
  
  void Mesh_::create_element_matrix_buffer(int is_schur) {
    assert(this->get_scheduler()->is_loaded());
    Scheduler s = this->get_scheduler();
    if (is_schur)
      s->execute_elements_seq( &(op_create_matrix_buffer_with_schur), true );
    else
      s->execute_elements_seq( &(op_create_matrix_buffer_without_schur), true );
  }
  
  void Mesh_::create_matrix_buffer() { this->create_matrix_buffer(false); }
  void Mesh_::create_matrix_buffer(int is_schur) {
    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
        
      assert(e->is_matrix_created());
        
      // default is create all buffers except for ABR
      // ABR is schur complement 
      for (int i=UHM_ATL;i<UHM_END;i++) 
        if (i!=UHM_ABR)
          e->get_matrix()->create_buffer(i);
        
      // when it has been requested, then allocte for schur
      if (is_schur) 
        e->get_matrix()->create_buffer(UHM_ABR);
    }
      
  }

  void Mesh_::free_matrix() {
    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
        
      if (e->is_matrix_created()) {
        delete e->get_matrix();
        e->set_matrix(nil_matrix);
      }
    }
  }

  void Mesh_::free_matrix_buffer() {
    
    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      
      assert(e->is_matrix_created());
      e->get_matrix()->free_buffer();
    }
  }

  void Mesh_::random_matrix()     { this->_random_matrix( false ); }
  void Mesh_::random_spd_matrix() { this->_random_matrix( true ); }
  void Mesh_::triangularize() {
    
    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      if ( e->is_leaf() && !e->is_matrix_reusable() )
        e->get_matrix()->triangularize(FLA_LOWER_TRIANGULAR);
    }
  }

  unsigned int Mesh_::get_n_dof() {
    unsigned int n_dof=0;
    std::map< int, Element_ >::iterator it;
    
    // collect the n_dof of factored nodes
    for (it=this->elements.begin();it!=this->elements.end();++it) {
      Element e = &(it->second);
      n_dof += e->get_n_dof().first;
    }
    return n_dof;
  }
  

  unsigned int Mesh_::get_n_nonzero_factor(int is_sym) {
    unsigned int n_nonzero=0;
    std::map< int, Element_ >::iterator it;
    
    // collect nonzeros which are in leaf UHM
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      std::pair< int, int > n_dof = e->get_n_dof();
      if (is_sym)
        n_nonzero += ( ( n_dof.first * n_dof.first  ) + 
                       ( n_dof.first * n_dof.second ) );
      else 
        n_nonzero += ( ( n_dof.first * n_dof.first  ) + 
                       ( n_dof.first * n_dof.second )*2 );
    }
    return n_nonzero;
  }

  unsigned int Mesh_::get_n_nonzero() {
    unsigned int n_nonzero=0;
    std::map< int, Element_ >::iterator it;

    // collect nonzeros which are in leaf UHM
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      if (e->is_leaf()) {
        std::pair< int, int > n_dof = e->get_n_dof();
        n_nonzero += ( (n_dof.first + n_dof.second)*
                       (n_dof.first + n_dof.second) );
      }
    }
    return n_nonzero;
  }

  void Mesh_::estimate_cost(int method, int datatype, int n_rhs,
                            double &flop_decompose, 
                            double &flop_solve,
                            double &buffer) {
    flop_decompose = 0.0;
    flop_solve     = 0.0;
    buffer         = 0.0;

    std::map< int, Element_ >::iterator it;

    // collect nonzeros which are in leaf UHM
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      double elt_decompose, elt_solve, elt_buffer;

      e->estimate_cost( method, datatype, n_rhs, 
                        elt_decompose, elt_solve, 
                        elt_buffer );

      flop_decompose   += elt_decompose;
      flop_solve       += elt_solve;
      buffer           += elt_buffer;

    }
  }

  void Mesh_::get_matrix_stat( int method, int datatype, 
                               int delta_dim,
                               std::vector< int > &x,
                               std::vector< int > &dist_uhm,
                               std::vector< double > &dist_flop,
                               std::vector< double > &dist_buffer) {
    int max_dim=0;
    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      std::pair<int,int> n_dof = e->get_n_dof();
      int dim = n_dof.first + n_dof.second;
      if (dim > max_dim) max_dim = dim;
    }

    int n_dist = max_dim/delta_dim + 1;
    x.clear();
    dist_uhm.clear();
    dist_flop.clear();
    dist_buffer.clear();
    for (int i=0;i<n_dist;++i) {
      x.push_back( i*delta_dim );
      dist_uhm.push_back   ( 0 );
      dist_flop.push_back  ( 0.0 );
      dist_buffer.push_back( 0.0 );
    }

    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      std::pair<int,int> n_dof = e->get_n_dof();
      int dim = n_dof.first + n_dof.second;

      int i = (dim/delta_dim);
      ++dist_uhm[i];

      double elt_decompose, elt_solve, elt_buffer;
      e->estimate_cost( method, datatype, 1,
                        elt_decompose, elt_solve,
                        elt_buffer );

      dist_flop[i] += elt_decompose;
      dist_buffer[i] += elt_buffer;
    }
  }

  double Mesh_::get_residual() {
    double rval = 0.0;
    std::map< int, Element_ >::iterator it;
#ifdef UHM_MULTITHREADING_ENABLE
    // ----------------------------------------------------------
    // ** UHM multi thread
    // ----------------------------------------------------------  
    // Use reduction to collect norm in each elements
    // and should NOT be inside task parallelism

    // change container from map to vector
    std::vector< Element > elts;
    elts.reserve(this->elements.size());
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      if (e->is_matrix_created()) elts.push_back(e);
    }

#pragma omp parallel for reduction(+:rval) schedule(static)
    for (int i=0;i<elts.size();++i) 
      rval += elts.at(i)->get_matrix()->get_residual();

#else
    // ----------------------------------------------------------
    // ** UHM single thread
    // ----------------------------------------------------------  
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      if (e->is_matrix_created()) 
        rval += (e->get_matrix()->get_residual());
    }
#endif

    return rval;
  }
 
  double Mesh_::get_lower_triangular_norm() {
    assert(this->get_scheduler()->is_loaded());
    Scheduler s = this->get_scheduler();

    double rval = 0.0;
    std::map< int, Element_ >::iterator it;
    // ----------------------------------------------------------
    // ** UHM single thread
    // ----------------------------------------------------------  
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      if (e->is_matrix_created()) 
        rval = max(rval, (e->get_matrix()->get_lower_triangular_norm()));
    }

    return rval;
  }

  void Mesh_::_random_matrix( int is_spd ) {

    std::map< int, Element_ >::iterator it;
    for (it=this->elements.begin();it!=this->elements.end();it++) {
      Element e = &(it->second);
      assert(e->is_matrix_created());
      if ( e->is_leaf() && !e->is_matrix_reusable() ) {
        e->get_matrix()->create_buffer();
        if (is_spd)
          e->get_matrix()->random();
        else
          e->get_matrix()->random_spd(FLA_LOWER_TRIANGULAR);
      }
    }
  }
}
